//
//  PLDirectory.m
//  Plankton
//
//  Created by Marc Schlichte on 2008-07-13.
//  Copyright 2008 Framework Labs. All rights reserved.
//

#import "PLDirectory.h"

#import "PLPlatform.h"

#define PLAP_TYPE @"_plap._udp"

@interface PLDirectory ()

@property(retain) NSMutableSet *searchResult;
@property(retain) NSMutableSet *outstandingPeers;
@property(copy) NSString *serviceToSearch;

- (void)noteAid:(NSString *)aid;
- (void)maybeFinishSearch;

@end


@implementation PLDirectory

@synthesize delegate, aid = dirAid;
@synthesize searchResult, outstandingPeers, serviceToSearch;
@synthesize discoverable, autoDiscover;

- (id)initWithPlatform:(PLPlatform *)platform_
{
    if (self = [super init]) {
        platform = [platform_ retain];
        serviceToAidsDict = [[NSMutableDictionary alloc] init];
        peers = [[NSMutableSet alloc] init];
        dirAid = [[platform createAidWithName:@"plankton-directory"] retain];
        
        queryPattern = [[PLPattern patternWithPredicateFormat:@"SELF[FIRST] = 'query'"] retain];
        [queryPattern setMessageTarget:self action:@selector(onQuery:)];
        
        [platform receive:queryPattern];
        
        ownService = [[NSNetService alloc] initWithDomain:@"" type:PLAP_TYPE name:dirAid port:[PLPlatform portOfAid:dirAid]];
        [ownService setDelegate:self];
        
        serviceBrowser = [[NSNetServiceBrowser alloc] init];
        [serviceBrowser setDelegate:self];
    }
    return self;
}

- (void)dealloc
{   
    [serviceBrowser release];
    [ownService release];
    
    [serviceToSearch release];
    [searchResult release];
    [outstandingPeers release];
    
    [queryPattern release];
    [dirAid release];
    [peers release];
    [serviceToAidsDict release];
    [platform release];
    
    [super dealloc];
}

- (void)registerAid:(NSString *)aid forService:(NSString *)service
{
    NSMutableArray *aids = [serviceToAidsDict objectForKey:service];
    if (!aids) {
        aids = [NSMutableArray array];
        [serviceToAidsDict setObject:aids forKey:service];
    }
    [aids addObject:aid];
}

- (void)unregisterAid:(NSString *)aid
{
    for (NSString *service in serviceToAidsDict) {
        NSMutableArray *aids = [serviceToAidsDict objectForKey:service];
        [aids removeObject:aid];
    }
}

- (BOOL)searchAidsForService:(NSString *)service
{
    // we only allow one concurrent search request at a time
    if (isSearching || !delegate) {
        return NO;
    }
    
    // initialize search state
    isSearching = YES;
    self.searchResult = [NSMutableSet set];
    self.outstandingPeers = [peers mutableCopy];
    self.serviceToSearch = service;
    
    // first search local directory
    NSMutableArray *aids = [serviceToAidsDict objectForKey:service];
    for (NSString *aid in aids) {
        [self noteAid:aid];
    }
    
    // next search peers concurrently
    for (NSString *peer in [peers copy]) {
        [platform send:[NSArray arrayWithObjects:@"query", service, self.aid, nil] 
                    to:peer];
    
        PLPattern *pattern = [PLPattern patternWithPredicateFormat:
                              @"SELF[0] = 'queryReply' AND SELF[1] = %@ AND SELF[3] = %@", 
                              peer, service];
        [pattern setMessageTarget:self action:@selector(onQueryReply:)];
        [pattern setTimeoutTarget:self action:@selector(onQueryReplyTimeout:) after:10.0];
        pattern.userInfo = peer;
        
        [platform receive:pattern];
    }
    
    // check if we are already finished in case of no peer
    [self maybeFinishSearch];
    
    return YES;
}

- (void)addPeer:(NSString *)aid
{
    [peers addObject:aid];
}

- (void)removePeer:(NSString *)aid
{
    [peers removeObject:aid];
}

- (void)onQuery:(id)message
{
    NSString *service = [message objectAtIndex:1];
    NSArray *aids = [serviceToAidsDict objectForKey:service];
    if (!aids) {
        aids = [NSArray array];
    }
    
    NSString *sender = [message objectAtIndex:2];
    [platform send:[NSArray arrayWithObjects:@"queryReply", self.aid, aids, service, nil] to:sender];
    
    [platform receive:queryPattern];
}

- (void)onQueryReply:(id)message
{
    NSString *peer = [message objectAtIndex:1];
    NSArray *aids = [message objectAtIndex:2];
    
    for (NSString *aid in aids) {
        [self noteAid:aid];
    }
    
    [outstandingPeers removeObject:peer];
    [self maybeFinishSearch];
}

- (void)onQueryReplyTimeout:(PLPattern *)pattern
{
    [outstandingPeers removeObject:pattern.userInfo];
    [self maybeFinishSearch];
}

- (void)noteAid:(NSString *)aid
{
    if (![searchResult member:aid]) {
        [searchResult addObject:aid];
        [delegate directory:self foundAid:aid forService:serviceToSearch];
    }
}

- (void)maybeFinishSearch
{
    // search is finished if number of outstanding peers is zero
    if ([outstandingPeers count] == 0) {
        isSearching = NO;
        [delegate directory:self searchDoneForService:serviceToSearch];
    }
}

- (void)setDiscoverable:(BOOL)val
{
    if (discoverable != val) {
        if (discoverable) {
            [ownService stop];
        }
        discoverable = val;
        if (discoverable) {
            [ownService publish];
        }
    }
}

- (void)setAutoDiscover:(BOOL)val
{
    if (autoDiscover != val) {
        if (autoDiscover) {
            [serviceBrowser stop];
        }
        autoDiscover = val;
        if (autoDiscover) {
            [serviceBrowser searchForServicesOfType:PLAP_TYPE inDomain:@""];
        }
    }
}

- (void)netServiceBrowser:(NSNetServiceBrowser *)browser
           didFindService:(NSNetService *)aNetService
               moreComing:(BOOL)moreComing
{
    [self addPeer:[aNetService name]];
}

- (void)netServiceBrowser:(NSNetServiceBrowser *)browser
         didRemoveService:(NSNetService *)aNetService
               moreComing:(BOOL)moreComing
{
    [self removePeer:[aNetService name]];
}

@end